"""
数据管理器 - 统一管理各种数据存储
提供统一的数据访问接口，协调向量存储、图谱存储和文档存储
"""

import asyncio
from typing import List, Dict, Any, Optional, Union

from ..core.interfaces import Document, Entity, Relationship, SearchResult
from ..core.events import EventEmitter
from ..core.config import config

from .vector_store.chroma_store import ChromaVectorStore
from .graph_store.neo4j_store import Neo4jGraphStore
from .document_store.file_store import FileDocumentStore


class DataManager(EventEmitter):
    """数据管理器 - 协调所有数据存储层"""
    
    def __init__(self):
        super().__init__()
        self.vector_store = None
        self.graph_store = None
        self.document_store = None
        self._initialized = False
    
    async def initialize(self):
        """初始化数据管理器"""
        if self._initialized:
            return
        
        try:
            await self.emit("data_manager_initialization_started", {})
            
            # 初始化向量存储
            self.vector_store = ChromaVectorStore(
                persist_directory="./data/vector_db",
                collection_name="hr_documents"
            )
            await self.vector_store.initialize()
            
            # 初始化文档存储
            self.document_store = FileDocumentStore("./data/documents")
            await self.document_store.initialize()
            
            # 初始化图谱存储（可选）
            try:
                neo4j_config = config.get("data_layer.graph_store", {})
                if neo4j_config.get("enabled", False):
                    self.graph_store = Neo4jGraphStore(
                        uri=neo4j_config.get("uri", "bolt://localhost:7687"),
                        username=neo4j_config.get("username", "neo4j"),
                        password=neo4j_config.get("password", "password")
                    )
                    await self.graph_store.initialize()
            except Exception as e:
                print(f"Warning: Neo4j not available: {e}")
                self.graph_store = None
            
            self._initialized = True
            
            await self.emit("data_manager_initialized", {
                "vector_store": "ChromaDB",
                "document_store": "FileStore",
                "graph_store": "Neo4j" if self.graph_store else "None"
            })
            
        except Exception as e:
            await self.emit_error("data_manager_initialization", e)
            raise
    
    async def store_document(
        self, 
        content: str, 
        metadata: Dict[str, Any] = None,
        generate_embedding: bool = True
    ) -> str:
        """存储文档到所有相关存储层"""
        try:
            if not self._initialized:
                await self.initialize()
            
            # 存储到文档存储
            doc_id = await self.document_store.store_document(content, metadata)
            
            # 创建文档对象
            document = Document(
                id=doc_id,
                content=content,
                metadata=metadata or {}
            )
            
            # 存储到向量存储
            if generate_embedding:
                await self.vector_store.add_documents([document])
            
            await self.emit("document_stored_all", {
                "doc_id": doc_id,
                "stores": ["document", "vector"]
            })
            
            return doc_id
            
        except Exception as e:
            await self.emit_error("document_store_all", e)
            raise
    
    async def search_documents(
        self,
        query: str,
        search_type: str = "hybrid",
        top_k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """多模态文档搜索"""
        try:
            if not self._initialized:
                await self.initialize()
            
            results = []
            
            # 向量搜索
            if search_type in ["vector", "hybrid"]:
                vector_results = await self.vector_store.search(query, top_k)
                results.extend(vector_results)
            
            # 关键词搜索（简单实现）
            if search_type in ["keyword", "hybrid"]:
                keyword_results = await self._keyword_search(query, top_k, filter_metadata)
                results.extend(keyword_results)
            
            # 图谱搜索
            if search_type in ["graph", "hybrid"] and self.graph_store:
                graph_results = await self._graph_search(query, top_k)
                results.extend(graph_results)
            
            # 去重和排序
            unique_results = self._deduplicate_results(results)
            sorted_results = sorted(unique_results, key=lambda x: x.score, reverse=True)
            
            return sorted_results[:top_k]
            
        except Exception as e:
            await self.emit_error("document_search", e)
            return []
    
    async def _keyword_search(
        self,
        query: str,
        top_k: int,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """简单的关键词搜索实现"""
        results = []
        
        # 获取所有文档
        doc_ids = await self.document_store.list_documents(filter_metadata)
        
        # 简单的关键词匹配
        for doc_id in doc_ids:
            document = await self.document_store.get_document(doc_id)
            if document and query.lower() in document.content.lower():
                # 简单评分：关键词出现次数
                score = document.content.lower().count(query.lower()) / len(document.content.split())
                results.append(SearchResult(
                    document=document,
                    score=score,
                    relevance_type="keyword"
                ))
        
        return sorted(results, key=lambda x: x.score, reverse=True)[:top_k]
    
    async def _graph_search(self, query: str, top_k: int) -> List[SearchResult]:
        """图谱搜索实现（简化版）"""
        if not self.graph_store:
            return []
        
        # 这里应该实现图谱搜索逻辑
        # 暂时返回空结果
        return []
    
    def _deduplicate_results(self, results: List[SearchResult]) -> List[SearchResult]:
        """去重搜索结果"""
        seen_docs = set()
        unique_results = []
        
        for result in results:
            if result.document.id not in seen_docs:
                seen_docs.add(result.document.id)
                unique_results.append(result)
        
        return unique_results
    
    async def add_entity(self, entity: Entity):
        """添加实体到图谱"""
        if self.graph_store:
            await self.graph_store.add_entity(entity)
    
    async def add_relationship(self, relationship: Relationship):
        """添加关系到图谱"""
        if self.graph_store:
            await self.graph_store.add_relationship(relationship)
    
    async def get_storage_stats(self) -> Dict[str, Any]:
        """获取存储统计信息"""
        stats = {}
        
        if self.vector_store:
            stats["vector_store"] = await self.vector_store.get_collection_stats()
        
        if self.document_store:
            stats["document_store"] = await self.document_store.get_storage_stats()
        
        if self.graph_store:
            stats["graph_store"] = await self.graph_store.get_graph_stats()
        
        return stats
    
    async def cleanup(self):
        """清理资源"""
        try:
            if self.graph_store:
                await self.graph_store.close()
            
            await self.emit("data_manager_cleanup_completed", {})
            
        except Exception as e:
            await self.emit_error("data_manager_cleanup", e)